#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>
#include <sys/socket.h>
#include <netinet/in.h>
#include <arpa/inet.h>
#include <time.h>
#include <signal.h>
#include <errno.h>
#include <pthread.h>
#include <sys/time.h>
#include <ifaddrs.h>
#include <net/if.h>
#ifdef HAVE_GETOPT_H
#include <getopt.h>
#else
// 简单的getopt实现
extern char *optarg;
extern int optind, opterr, optopt;
int getopt(int argc, char * const argv[], const char *optstring);
#endif

#define BUFFER_SIZE 1024 * 1024  // 1MB buffer
#define DEFAULT_PORT 8888
#define DEFAULT_SIZE 100  // 默认发送100MB
#define DEFAULT_THREADS 4  // 默认4个并发连接
#define MAX_IPS 16  // 最大支持的IP数量

// 全局变量
volatile int running = 1;
pthread_mutex_t stats_mutex = PTHREAD_MUTEX_INITIALIZER;
long total_bytes = 0;
time_t start_time, end_time;

// IP地址信息结构
typedef struct {
    char ip_str[INET_ADDRSTRLEN];
    struct in_addr ip_addr;
    int interface_index;
} ip_info_t;

// 线程参数结构
typedef struct {
    int thread_id;
    char *server_ip;
    int port;
    long target_bytes;
    int test_time_sec;
    int use_time_mode;
    long thread_bytes;
    int client_fd;
    ip_info_t *source_ip;  // 绑定的源IP
    int source_port;       // 绑定的源端口
} thread_args_t;

void signal_handler(int sig) {
    (void)sig;  // 避免未使用参数警告
    printf("\n正在停止传输...\n");
    running = 0;
}

void print_usage(const char *program_name) {
    printf("用法: %s [选项]\n", program_name);
    printf("选项:\n");
    printf("  -s <服务器IP>  指定服务器IP地址 (默认: 127.0.0.1)\n");
    printf("  -p <端口>      指定服务器端口 (默认: %d)\n", DEFAULT_PORT);
    printf("  -d <大小>      指定发送数据大小(MB) (默认: %d)\n", DEFAULT_SIZE);
    printf("  -t <时间>      指定测试时间(秒) (与-d互斥)\n");
    printf("  -n <线程数>    指定并发连接数 (默认: %d)\n", DEFAULT_THREADS);
    printf("  -i <网卡>      指定要使用的网络接口 (如: eth0, bond0)\n");
    printf("  -m             启用多IP绑定模式 (自动发现并使用多个IP)\n");
    printf("  -h             显示此帮助信息\n");
    printf("\n示例:\n");
    printf("  %s -s 192.168.1.100 -d 500 -n 8 -m    # 多IP绑定测试\n", program_name);
    printf("  %s -s 192.168.1.100 -i bond0 -n 16    # 指定bond接口\n", program_name);
    printf("\n注意: 使用-m选项可以更好地测试bond网络的聚合带宽\n");
}

// 获取网络接口的IP地址
int get_interface_ips(const char *interface_name, ip_info_t *ips, int max_ips) {
    struct ifaddrs *ifaddr, *ifa;
    int count = 0;
    
    if (getifaddrs(&ifaddr) == -1) {
        perror("getifaddrs");
        return 0;
    }
    
    for (ifa = ifaddr; ifa != NULL && count < max_ips; ifa = ifa->ifa_next) {
        if (ifa->ifa_addr == NULL) continue;
        
        // 只处理IPv4地址
        if (ifa->ifa_addr->sa_family != AF_INET) continue;
        
        // 如果指定了接口名，只获取该接口的IP
        if (interface_name && strcmp(ifa->ifa_name, interface_name) != 0) continue;
        
        // 跳过回环接口
        if (strcmp(ifa->ifa_name, "lo") == 0) continue;
        
        struct sockaddr_in *addr_in = (struct sockaddr_in *)ifa->ifa_addr;
        
        // 检查接口是否活跃
        if (!(ifa->ifa_flags & IFF_UP)) continue;
        
        strcpy(ips[count].ip_str, inet_ntoa(addr_in->sin_addr));
        ips[count].ip_addr = addr_in->sin_addr;
        ips[count].interface_index = if_nametoindex(ifa->ifa_name);
        
        printf("发现IP: %s (接口: %s)\n", ips[count].ip_str, ifa->ifa_name);
        count++;
    }
    
    freeifaddrs(ifaddr);
    return count;
}

// 绑定socket到指定的IP和端口
int bind_socket_to_ip_port(int sockfd, const char *ip_str, int port) {
    struct sockaddr_in local_addr;
    
    memset(&local_addr, 0, sizeof(local_addr));
    local_addr.sin_family = AF_INET;
    local_addr.sin_port = htons(port);
    
    if (inet_pton(AF_INET, ip_str, &local_addr.sin_addr) <= 0) {
        return -1;
    }
    
    if (bind(sockfd, (struct sockaddr *)&local_addr, sizeof(local_addr)) == -1) {
        return -1;
    }
    
    return 0;
}

// 工作线程函数
void* worker_thread(void* arg) {
    thread_args_t *args = (thread_args_t*)arg;
    char buffer[BUFFER_SIZE];
    ssize_t bytes_sent;
    long thread_total = 0;
    
    // 初始化发送缓冲区
    memset(buffer, 'A', BUFFER_SIZE);
    
    // 创建socket
    args->client_fd = socket(AF_INET, SOCK_STREAM, 0);
    if (args->client_fd == -1) {
        perror("创建socket失败");
        return NULL;
    }
    
    // 如果指定了源IP，绑定到该IP
    if (args->source_ip) {
        if (bind_socket_to_ip_port(args->client_fd, args->source_ip->ip_str, args->source_port) == -1) {
            printf("线程%d: 绑定到 %s:%d 失败，使用系统默认\n", 
                   args->thread_id, args->source_ip->ip_str, args->source_port);
        } else {
            printf("线程%d: 绑定到 %s:%d\n", 
                   args->thread_id, args->source_ip->ip_str, args->source_port);
        }
    }
    
    // 配置服务器地址
    struct sockaddr_in server_addr;
    memset(&server_addr, 0, sizeof(server_addr));
    server_addr.sin_family = AF_INET;
    server_addr.sin_port = htons(args->port);
    
    if (inet_pton(AF_INET, args->server_ip, &server_addr.sin_addr) <= 0) {
        fprintf(stderr, "线程%d: 无效的IP地址: %s\n", args->thread_id, args->server_ip);
        close(args->client_fd);
        return NULL;
    }
    
    // 连接到服务器
    if (connect(args->client_fd, (struct sockaddr *)&server_addr, sizeof(server_addr)) == -1) {
        perror("连接服务器失败");
        close(args->client_fd);
        return NULL;
    }
    
    printf("线程%d: 连接成功\n", args->thread_id);
    
    // 开始传输数据
    while (running) {
        // 检查是否达到目标
        if (!args->use_time_mode && thread_total >= args->target_bytes) {
            break;
        }
        if (args->use_time_mode && (time(NULL) - start_time) >= args->test_time_sec) {
            break;
        }
        
        // 发送数据
        bytes_sent = send(args->client_fd, buffer, BUFFER_SIZE, 0);
        if (bytes_sent == -1) {
            if (errno == EINTR) {
                continue;  // 被信号中断，继续循环
            }
            perror("发送数据失败");
            break;
        }
        
        thread_total += bytes_sent;
        
        // 更新全局统计
        pthread_mutex_lock(&stats_mutex);
        total_bytes += bytes_sent;
        pthread_mutex_unlock(&stats_mutex);
    }
    
    args->thread_bytes = thread_total;
    close(args->client_fd);
    
    printf("线程%d: 传输完成，发送 %.2f MB", 
           args->thread_id, thread_total / (1024.0 * 1024.0));
    if (args->source_ip) {
        printf(" (源IP: %s)", args->source_ip->ip_str);
    }
    printf("\n");
    
    return NULL;
}

// 统计显示线程
void* stats_thread(void* arg) {
    thread_args_t *args = (thread_args_t*)arg;
    time_t last_update_time = start_time;
    
    while (running) {
        sleep(1);
        
        time_t current_time = time(NULL);
        if (current_time != last_update_time) {
            double duration = difftime(current_time, start_time);
            if (duration > 0) {
                pthread_mutex_lock(&stats_mutex);
                long current_total = total_bytes;
                pthread_mutex_unlock(&stats_mutex);
                
                double bandwidth_mbps = (current_total * 8.0) / (duration * 1024 * 1024);
                
                if (args->use_time_mode) {
                    printf("\r总发送: %.2f MB, 聚合带宽: %.2f Mbps, 剩余时间: %ld 秒", 
                           current_total / (1024.0 * 1024.0), bandwidth_mbps, 
                           args->test_time_sec - (current_time - start_time));
                } else {
                    printf("\r总发送: %.2f MB, 聚合带宽: %.2f Mbps, 进度: %.1f%%", 
                           current_total / (1024.0 * 1024.0), bandwidth_mbps,
                           (double)current_total / (args->target_bytes * args->thread_id) * 100);
                }
                fflush(stdout);
            }
            last_update_time = current_time;
        }
    }
    
    return NULL;
}

int main(int argc, char *argv[]) {
    char *server_ip = "127.0.0.1";
    int port = DEFAULT_PORT;
    int data_size_mb = DEFAULT_SIZE;
    int test_time_sec = 0;
    int num_threads = DEFAULT_THREADS;
    char *interface_name = NULL;
    int multi_ip_mode = 0;
    int opt;
    int use_time_mode = 0;
    pthread_t *threads;
    pthread_t stats_tid;
    thread_args_t *thread_args;
    ip_info_t available_ips[MAX_IPS];
    int ip_count = 0;
    int i;
    double duration, bandwidth_mbps;
    
    // 解析命令行参数
    while ((opt = getopt(argc, argv, "s:p:d:t:n:i:mh")) != -1) {
        switch (opt) {
            case 's':
                server_ip = optarg;
                break;
            case 'p':
                port = atoi(optarg);
                if (port <= 0 || port > 65535) {
                    fprintf(stderr, "错误: 端口号必须在1-65535之间\n");
                    exit(EXIT_FAILURE);
                }
                break;
            case 'd':
                data_size_mb = atoi(optarg);
                if (data_size_mb <= 0) {
                    fprintf(stderr, "错误: 数据大小必须大于0\n");
                    exit(EXIT_FAILURE);
                }
                break;
            case 't':
                test_time_sec = atoi(optarg);
                if (test_time_sec <= 0) {
                    fprintf(stderr, "错误: 测试时间必须大于0\n");
                    exit(EXIT_FAILURE);
                }
                use_time_mode = 1;
                break;
            case 'n':
                num_threads = atoi(optarg);
                if (num_threads <= 0 || num_threads > 64) {
                    fprintf(stderr, "错误: 线程数必须在1-64之间\n");
                    exit(EXIT_FAILURE);
                }
                break;
            case 'i':
                interface_name = optarg;
                break;
            case 'm':
                multi_ip_mode = 1;
                break;
            case 'h':
                print_usage(argv[0]);
                exit(EXIT_SUCCESS);
            default:
                print_usage(argv[0]);
                exit(EXIT_FAILURE);
        }
    }
    
    // 检查参数冲突
    if (use_time_mode && data_size_mb != DEFAULT_SIZE) {
        fprintf(stderr, "错误: -d和-t参数不能同时使用\n");
        exit(EXIT_FAILURE);
    }
    
    // 设置信号处理
    signal(SIGINT, signal_handler);
    signal(SIGTERM, signal_handler);
    
    // 获取可用的IP地址
    if (multi_ip_mode || interface_name) {
        printf("扫描网络接口IP地址...\n");
        ip_count = get_interface_ips(interface_name, available_ips, MAX_IPS);
        if (ip_count == 0) {
            fprintf(stderr, "错误: 未找到可用的IP地址\n");
            exit(EXIT_FAILURE);
        }
        printf("找到 %d 个可用IP地址\n", ip_count);
    }
    
    // 分配内存
    threads = malloc(num_threads * sizeof(pthread_t));
    thread_args = malloc(num_threads * sizeof(thread_args_t));
    if (!threads || !thread_args) {
        fprintf(stderr, "错误: 内存分配失败\n");
        exit(EXIT_FAILURE);
    }
    
    printf("增强版TCP带宽测试客户端 (支持bond网络)\n");
    printf("========================================\n");
    printf("服务器: %s:%d\n", server_ip, port);
    printf("并发连接数: %d\n", num_threads);
    if (interface_name) {
        printf("指定接口: %s\n", interface_name);
    }
    if (multi_ip_mode) {
        printf("多IP模式: 启用 (使用 %d 个IP)\n", ip_count);
    }
    if (use_time_mode) {
        printf("测试时间: %d 秒\n", test_time_sec);
    } else {
        printf("总数据大小: %d MB (每个连接: %.2f MB)\n", 
               data_size_mb, (double)data_size_mb / num_threads);
    }
    printf("按Ctrl+C停止传输\n\n");
    
    // 初始化线程参数
    for (i = 0; i < num_threads; i++) {
        thread_args[i].thread_id = i + 1;
        thread_args[i].server_ip = server_ip;
        thread_args[i].port = port;
        thread_args[i].use_time_mode = use_time_mode;
        thread_args[i].test_time_sec = test_time_sec;
        thread_args[i].target_bytes = use_time_mode ? 0 : (long)data_size_mb * 1024 * 1024 / num_threads;
        thread_args[i].thread_bytes = 0;
        thread_args[i].client_fd = -1;
        
        // 分配源IP和端口
        if (multi_ip_mode && ip_count > 0) {
            thread_args[i].source_ip = &available_ips[i % ip_count];
            thread_args[i].source_port = 10000 + i;  // 使用不同的源端口
        } else {
            thread_args[i].source_ip = NULL;
            thread_args[i].source_port = 0;
        }
    }
    
    // 记录开始时间
    start_time = time(NULL);
    
    // 创建统计显示线程
    if (pthread_create(&stats_tid, NULL, stats_thread, &thread_args[0]) != 0) {
        perror("创建统计线程失败");
        free(threads);
        free(thread_args);
        exit(EXIT_FAILURE);
    }
    
    // 创建工作线程
    for (i = 0; i < num_threads; i++) {
        if (pthread_create(&threads[i], NULL, worker_thread, &thread_args[i]) != 0) {
            perror("创建工作线程失败");
            running = 0;
            break;
        }
    }
    
    // 等待所有工作线程完成
    for (i = 0; i < num_threads; i++) {
        pthread_join(threads[i], NULL);
    }
    
    // 停止统计线程
    running = 0;
    pthread_join(stats_tid, NULL);
    
    // 计算最终统计信息
    end_time = time(NULL);
    duration = difftime(end_time, start_time);
    
    printf("\n\n=== 聚合带宽测试结果 ===\n");
    printf("并发连接数: %d\n", num_threads);
    if (multi_ip_mode) {
        printf("使用IP数量: %d\n", ip_count);
    }
    printf("总发送数据: %.2f MB\n", total_bytes / (1024.0 * 1024.0));
    printf("传输时间: %.2f 秒\n", duration);
    if (duration > 0) {
        bandwidth_mbps = (total_bytes * 8.0) / (duration * 1024 * 1024);
        printf("聚合带宽: %.2f Mbps\n", bandwidth_mbps);
        printf("平均速度: %.2f MB/s\n", total_bytes / (duration * 1024 * 1024));
        printf("单连接平均带宽: %.2f Mbps\n", bandwidth_mbps / num_threads);
    }
    
    // 显示每个线程的统计
    printf("\n各线程统计:\n");
    for (i = 0; i < num_threads; i++) {
        printf("  线程%d: %.2f MB", i + 1, thread_args[i].thread_bytes / (1024.0 * 1024.0));
        if (thread_args[i].source_ip) {
            printf(" (源IP: %s)", thread_args[i].source_ip->ip_str);
        }
        printf("\n");
    }
    printf("========================\n");
    
    // 清理资源
    free(threads);
    free(thread_args);
    pthread_mutex_destroy(&stats_mutex);
    
    return 0;
}
